/*
 * Copyright 2023 Salesforce, Inc. All rights reserved.
 * The software in this package is published under the terms of the CPAL v1.0
 * license, a copy of which has been included with this distribution in the
 * LICENSE.txt file.
 */
package org.mule.runtime.core.privileged.routing;

import static java.util.Collections.singletonList;
import static java.util.stream.Collectors.toList;
import static org.mule.runtime.core.privileged.event.PrivilegedEvent.setCurrentEvent;

import org.mule.runtime.api.message.Message;
import org.mule.runtime.core.api.event.CoreEvent;
import org.mule.runtime.core.internal.message.InternalMessage;
import org.mule.runtime.core.privileged.event.PrivilegedEvent;

import java.util.ArrayList;
import java.util.List;

/**
 * The default results handler for all outbound endpoint. Depending on the number of messages passed it the returning message will
 * be different. If the 'results' param is null or empty, null is returned. If the 'results' param contains a single
 * {@link InternalMessage}, than that message is returned.
 * <p/>
 * Note that right now (as of Mule 2.0.1) this SPI is not pluggable and this implementation is the default and only
 * implementation.
 *
 * @see InternalMessage
 * @see org.mule.runtime.core.DefaultMessageCollection
 */
public final class DefaultRouterResultsHandler implements RouterResultsHandler {

  private boolean returnCollectionWithSingleResult = false;

  public DefaultRouterResultsHandler() {}

  /**
   * @param returnCollectionWithSingleResult if a MuleMessageCollection should be return despite there's only one result event
   */
  public DefaultRouterResultsHandler(boolean returnCollectionWithSingleResult) {
    this.returnCollectionWithSingleResult = returnCollectionWithSingleResult;
  }

  /**
   * Aggregates the events in the results list into one single {@link CoreEvent} You should only use this method when you're sure
   * that all the events in the results list were generated by the same thread that's going to handle the aggregated response
   *
   * @param results
   * @param previous
   * @return
   */
  @Override
  @SuppressWarnings(value = {"unchecked"})
  public CoreEvent aggregateResults(final List<CoreEvent> results,
                                    final CoreEvent previous) {
    if (results == null) {
      return null;
    } else if (results.size() == 1) {
      PrivilegedEvent event = (PrivilegedEvent) results.get(0);
      if (event == null) {
        return event;
      } else if (event != null && event.getMessage() != null) {
        if (returnCollectionWithSingleResult) {
          return createMessageCollectionWithSingleMessage(event);
        } else {
          return event;
        }
      } else {
        return previous;
      }
    } else {
      List<CoreEvent> nonNullResults = results.stream().filter(object -> object != null &&
          object.getMessage() != null).collect(toList());

      if (nonNullResults.size() == 0) {
        return CoreEvent.builder(previous).message(Message.of(null)).build();
      } else if (nonNullResults.size() == 1) {
        return nonNullResults.get(0);
      } else {
        return createMessageCollection(nonNullResults, previous);
      }
    }
  }

  private CoreEvent createMessageCollectionWithSingleMessage(PrivilegedEvent event) {
    final Message coll = Message.builder().collectionValue(singletonList(event.getMessage()), Message.class).build();
    event = PrivilegedEvent.builder(event).message(coll).build();
    setCurrentEvent(event);
    return event;
  }

  private CoreEvent createMessageCollection(final List<CoreEvent> nonNullResults,
                                            final CoreEvent previous) {
    final PrivilegedEvent.Builder resultBuilder = PrivilegedEvent.builder(previous);

    List<Message> list = new ArrayList<>();
    for (CoreEvent event : nonNullResults) {
      for (String flowVarName : event.getVariables().keySet()) {
        resultBuilder.addVariable(flowVarName, event.getVariables().get(flowVarName));
      }
      list.add(event.getMessage());
    }
    final Message coll = Message.builder().collectionValue(list, Message.class).build();

    PrivilegedEvent resultEvent = resultBuilder.message(coll).build();
    setCurrentEvent(resultEvent);
    return resultEvent;
  }
}
